import numpy as np
import matplotlib.pyplot as plt
from sklearn.metrics import silhouette_score, calinski_harabasz_score, davies_bouldin_score

class KMeansCluster:
    def __init__(self, k=3, max_iters=100, plot_steps=False):
        self.k = k
        self.max_iters = max_iters
        self.plot_steps = plot_steps
        
        # 初始化簇心
        self.centroids = None
        # 标签列表
        self.clusters = [[] for _ in range(self.k)]

    def _initialize_centroids(self, X):
        """随机选择 k 个样本作为初始质心"""
        idx = np.random.choice(len(X), self.k, replace=False)
        self.centroids = X[idx]

    def _closest_centroid(self, sample):
        """计算样本到所有质心的距离，返回最近的那个"""
        distances = [np.linalg.norm(sample - point) for point in self.centroids]
        closest = np.argmin(distances)
        return closest

    def _create_clusters(self, X):
        """分配每个样本给最近的质心"""
        self.clusters = [[] for _ in range(self.k)]
        for idx, sample in enumerate(X):
            centroid_idx = self._closest_centroid(sample)
            self.clusters[centroid_idx].append(idx)

    def _calculate_new_centroids(self, X):
        """重新计算新的质心位置"""
        centroids = np.zeros((self.k, X.shape[1]))
        for i, cluster in enumerate(self.clusters):
            if cluster:  # 确保簇不为空
                new_centroid = np.mean(X[cluster], axis=0)
                centroids[i] = new_centroid
        return centroids

    def predict(self, X,SHOW):
        """为每个样本分配一个标签"""
        self._initialize_centroids(X)
        metrics = {}
        for it in range(self.max_iters):
            self._create_clusters(X)
            previous_centroids = self.centroids
            self.centroids = self._calculate_new_centroids(X)
            
            # 每次迭代时计算并打印聚类质量指标
            labels = np.empty(len(X))
            for cluster_idx, cluster in enumerate(self.clusters):
                for sample_idx in cluster:
                    labels[sample_idx] = cluster_idx
            
            sse = round(self.calculate_sse(X, labels),2)
            silhouette = round(silhouette_score(X, labels),2)
            ch_index = round(calinski_harabasz_score(X, labels),2)
            db_index = round(davies_bouldin_score(X, labels),2)
            iteration = it+1
            if SHOW:
                print(f"迭代 {iteration}:")
                print(f"  - 误差平方和 (SSE): {sse:.2f}")
                print(f"  - 轮廓系数 (Silhouette Score): {silhouette:.2f}")
                print(f"  - CH 指数 (Calinski-Harabasz Index): {ch_index:.2f}")
                print(f"  - DB 指数 (Davies-Bouldin Index): {db_index:.2f}")
                print("=" * 50)
            metrics[iteration] = {
                "SSE": sse,
                "Silhouette Score": silhouette,
                "Calinski-Harabasz Index": ch_index,
                "Davies-Bouldin Index": db_index
        }
            if self.plot_steps:
                self.plot(X, it)
                
            if np.allclose(previous_centroids, self.centroids):
                break
        
        # 最终标签分配
        labels = np.empty(len(X))
        for cluster_idx, cluster in enumerate(self.clusters):
            for sample_idx in cluster:
                labels[sample_idx] = cluster_idx
        last_metrics = metrics[iteration]

        return labels,last_metrics,iteration

    def calculate_sse(self, X, labels):
        """计算误差平方和 (SSE)"""
        sse = 0
        for i, cluster in enumerate(self.clusters):
            if cluster:
                sse += np.sum((X[cluster] - self.centroids[i]) ** 2)
        return sse

    def plot(self, X, iteration):
        fig, ax = plt.subplots(figsize=(12, 8))
        colors = ['r', 'g', 'b', 'y', 'c', 'm']
        for i, index in enumerate(self.clusters):
            point = X[index].T
            ax.scatter(*point, color=colors[i % len(colors)], label=f'Cluster {i+1}')
            
        for point in self.centroids:
            ax.scatter(*point, marker='x', color='black', linewidth=2)
            
        plt.xlabel('Feature_1')
        plt.ylabel('Feature_2')
        plt.title(f'Iteration {iteration+1}')
        plt.legend()
        plt.show()
    def plot_metrics(self, data, class_metrics, metric_name):
        fig, ax = plt.subplots(figsize=(12, 8))
        colors = ['r', 'g', 'b', 'y', 'c', 'm']
        for i, index in enumerate(class_metrics.clusters):
            point = data[index].T
            ax.scatter(*point, color=colors[i % len(colors)], label=f'Cluster {i+1}')
            
        for point in class_metrics.centroids:
            ax.scatter(*point, marker='x', color='black', linewidth=2)
        plt.title(f'{metric_name}')